<template>
  <div class="container">
    <div class="main-box">
      <div class="header-line">
        <div class="title">
          <img src="../assets/image/param.png" alt="" class="title-icon">
          超参数设置
        </div>
        <div class="operate-btn">
          <el-button type="info" @click="configForm = defaultConfig" size="mini">重置为默认</el-button>
          <el-button type="primary" size="mini" @click="saveConfig">保存并使用</el-button>
        </div>
      </div>
      <el-form ref="form" :model="configForm" label-width="auto" style="margin-top: 10px">
        <el-row :gutter="24">
          <el-col :span="6">
            <el-form-item label="模型权重">
              <el-select v-model="configForm.weightName" placeholder="请选择模型权重">
                <el-option
                    v-for="item in weightOptions"
                    :key="item.value"
                    :label="item.label"
                    :value="item.value">
                </el-option>
              </el-select>
            </el-form-item>
          </el-col>
          <el-col :span="6">
            <el-form-item label="学习率">
              <el-input-number v-model="configForm.learningRate" :precision="4" :min="0" :max="1"
                               placeholder="请输入学习率大小"
                               :step="0.0001"></el-input-number>
            </el-form-item>
          </el-col>
          <el-col :span="6">
            <el-form-item label="训练轮数">
              <el-input v-model="configForm.epoch" placeholder="请输入训练轮数"></el-input>
            </el-form-item>
          </el-col>
          <el-col :span="6">
            <el-form-item label="多分类阈值">
              <el-input v-model="configForm.threshold" placeholder="请输入多分类阈值"></el-input>
            </el-form-item>
          </el-col>
        </el-row>
      </el-form>
      <div class="header-line">
        <div class="title">
          <img src="../assets/image/table.png" alt="" class="title-icon">
          训练/预测结果
        </div>
        <div class="operate-btn">
          <el-button type="success" @click="modelStartWork(0)" size="mini">开始训练</el-button>
          <el-button type="warning" @click="modelStartWork(1)" size="mini">开始预测</el-button>
          <el-button type="danger" @click="stopModel" size="mini">结束运行</el-button>
        </div>
      </div>
      <el-carousel height="600px" v-if="allResultTable.length > 0">
        <el-carousel-item v-for="(item,index) in allResultTable" :key="index">
          <div class="table-title">{{ item.title }}</div>
          <el-table
              :data="item.table"
              border
              stripe
              style="width: 100%;margin-top: 10px">
            <el-table-column
                prop="disease"
                label="疾病名称"
                width="180">
            </el-table-column>
            <el-table-column
                prop="accuracy"
                label="准确率(ACC)">
            </el-table-column>
            <el-table-column
                prop="precision"
                label="查准率(P)">
            </el-table-column>
            <el-table-column
                prop="recall"
                label="召回率(R)">
            </el-table-column>
            <el-table-column
                prop="f1"
                label="F1">
            </el-table-column>
            <el-table-column
                prop="AUC"
                label="AUC">
            </el-table-column>
          </el-table>
        </el-carousel-item>
      </el-carousel>
      <div v-else>
        <el-table
            :data="exampleResultTable"
            border
            stripe
            style="width: 100%;margin-top: 10px">
          <el-table-column
              prop="disease"
              label="疾病名称"
              width="180">
          </el-table-column>
          <el-table-column
              prop="accuracy"
              label="准确率(ACC)">
          </el-table-column>
          <el-table-column
              prop="precision"
              label="查准率(P)">
          </el-table-column>
          <el-table-column
              prop="recall"
              label="召回率(R)">
          </el-table-column>
          <el-table-column
              prop="f1"
              label="F1">
          </el-table-column>
          <el-table-column
              prop="AUC"
              label="AUC">
          </el-table-column>
        </el-table>
      </div>
    </div>
  </div>
</template>

<script>
import {io} from "socket.io-client";
import {Message} from "element-ui";
import {getAllWeights, getDefaultConfig, saveAndUseConfig} from "@/api/diagnosis";

export default {
  name: 'ModelManagement',
  data() {
    return {
      // 默认配置
      defaultConfig: {},
      socket: null,
      // 每次数据集训练完一个batch后返回 记住是第几轮 已经是什么模式
      trainCount: 0,
      validCount: 0,
      mode: '',
      progressMessage: null,
      // 全部权重列表
      weightOptions: [],
      configForm: {
        doTrain: false,
        doTest: false,
        weightName: '',
        learningRate: '',
        epoch: '',
        threshold: ''
      },
      configFormRules: {
        weightName: [
          {required: true, message: '请选择模型权重', trigger: 'change'}
        ],
        learningRate: [
          {required: true, message: '请输入学习率', trigger: 'blur'}
        ],
        threshold: [
          {required: true, message: '请输入阈值', trigger: 'change'}
        ],
        epoch: [
          {required: true, message: '请输入训练轮数', trigger: 'blur'}
        ],
      },
      allResultTable: [],
      exampleResultTable: [
        {
          disease: 'N',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'D',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'G',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'C',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'A',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'H',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'M',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: 'O',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        },
        {
          disease: '平均值',
          accuracy: 0,
          precision: 0,
          recall: 0,
          f1: 0,
          AUC: 0
        }]
    }
  },
  mounted() {
    this.getWeightList();
    getDefaultConfig().then(res => {
      this.defaultConfig = res.data;
      this.defaultConfig['doTrain'] = false;
      this.defaultConfig['doTest'] = false;
    })
  },
  beforeDestroy() {
    // 离开时断开Socketio连接
    if (this.socket === null) {
      this.stopModel();
    }
  },
  methods: {
    // 保存并使用模型
    saveConfig() {
      saveAndUseConfig(this.configForm.weightName,this.configForm.learningRate,this.configForm.epoch,this.configForm.threshold).then(res => {
        this.$message.success('保存成功')
      })
    },
    // 获取全部权重
    getWeightList() {
      getAllWeights().then(res => {
        for (let i = 0; i < res.data.length; i++) {
          this.weightOptions.push({
            value: res.data[i],
            label: res.data[i]
          })
        }
      })
    },
    // 停止模型运行
    stopModel() {
      if (this.socket) {
        this.socket.emit('modelStop')
        this.socket.disconnect();
      }
    },
    modelStartWork(mode) {
      // 验证表单是否填写完整
      if (this.configForm.weightName === '') {
        this.$message.warning('请选择模型权重')
        return
      }
      if (this.configForm.learningRate === '') {
        this.$message.warning('请输入学习率')
        return
      }
      if (this.configForm.epoch === '') {
        this.$message.warning('请输入训练轮数')
        return
      }
      if (this.configForm.threshold === '') {
        this.$message.warning('请输入阈值')
      }
      if (mode === 0) {
        this.configForm.doTrain = true;
      }
      if (mode === 1) {
        this.configForm.doTest = true;
      }

      let loadingText = '模型正在运行，请稍后'
      let loading = this.$loading({
        lock: true,
        text: loadingText,
        spinner: 'el-icon-loading',
        background: 'rgba(0, 0, 0, 0.7)'
      })

      this.trainCount = 0;
      this.validCount = 0;

      this.socket = io("https://yiyuy.mynatapp.cc/model", {
        transports: ['websocket']
      })

      // 监听前端是否与后端socketio连接成功
      this.socket.on('connect', () => {
        loadingText = '数据集加载中...'
        loading.setText(loadingText)
        this.allResultTable = []
        this.socket.emit('receiveParam', {'configForm': this.configForm})
      })

      // 监听服务器那边返回的消息
      // 1. 模型每次返回的指标结果
      this.socket.on('model_response', (message) => {
        let scores = message.data.scores
        let resultTable = this.exampleResultTable
        for (let i = 0; i < resultTable.length - 1; i++) {
          resultTable[i].accuracy = scores['Accuracy Rate per class'][i] * 100
          resultTable[i].precision = scores['Precision per class'][i] * 100
          resultTable[i].recall = scores['Recall per class'][i] * 100
          resultTable[i].F1 = scores['F1-score per class'][i] * 100
          resultTable[i].AUC = scores['AUC per class'][i] * 100
        }

        let len = resultTable.length
        resultTable[len - 1].accuracy = message.data.acc * 100
        resultTable[len - 1].precision = message.data.precision * 100
        resultTable[len - 1].recall = message.data.recall * 100
        resultTable[len - 1].f1 = message.data.f1 * 100
        resultTable[len - 1].AUC = scores['Average AUC'] * 100

        this.progressMessage.close()
        this.progressMessage = null

        // 构造title
        let title = ''
        if ( this.mode === 'Train' ) {
          title = `第${this.trainCount}轮训练结果`
        }
        if ( this.mode === 'Valid' ) {
          title = `第${this.validCount}轮验证结果`
        }
        if ( this.mode === 'Predict' ) {
          title = `预测结果`
        }
        this.allResultTable.push({
          title: title,
          table: resultTable
        })
      })

      // 2.模型加载数据集结束后的反馈
      this.socket.on('dataset_end', (message) => {
        this.$message.success('数据集处理完毕，模型开始工作...')
        loading.close();
      })

      // 3.模型加载完一个batch的反馈
      this.socket.on('model_progress', (message) => {
        this.mode = message.mode;
        if (this.progressMessage === null) {
          let nowEpoch = 0
          this.mode === 'Train' ? nowEpoch = this.trainCount : nowEpoch = this.validCount
          this.mode === 'Train' ? this.trainCount ++ : nowEpoch = this.validCount ++
          this.progressMessage = Message({
            dangerouslyUseHTMLString: true,
            type: 'warning',
            message: `<strong><i>${this.mode} Epoch ${nowEpoch}</i></strong>： <progress type="warning" max="100" value="0" style="width: 250px" class="progress"></progress>`,
            duration: 0
          });
        }
        this.count += 1;
        let progressDom = document.querySelector('.progress')
        progressDom.setAttribute('value', (message.data * 100).toString())
      });

      // 4.模型训练完毕的反馈
      this.socket.on('model_end', (message) => {
        if (message.data === '[over]') {
          this.$message.success('模型运行完毕')
        }
        if (message.data === '[early over]') {
          this.$message.warning('模型运行提前结束')
        }
        this.stopModel()
      })

      // 监听当socketio出错时 反馈
      this.socket.on('error', (msg) => {
        console.log(msg)
        this.$message.error('服务器开小差了。')
        this.stopModel()
      });

    }
  }
}
</script>

<style lang="less" scoped>
@keyframes MoveIn {
  0% {
    opacity: 0;
    transform: translateY(12px);
  }

  100% {
    opacity: 1;
    transform: translateY(0);
  }
}

.container {
  width: 100%;
  height: 100%;
  animation: MoveIn .7s cubic-bezier(.4, 0, .2, 1);

  .main-box {
    width: 99%;
    height: 100%;
    margin: 0 auto;

    .header-line {
      width: 100%;
      display: flex;
      align-items: center;
      margin-top: 10px;

      .title {
        display: flex;
        align-items: center;
        font-weight: bold;
        color: #303133;

        .title-icon {
          width: 20px;
          height: 20px;
        }
      }

      .operate-btn {
        display: flex;
        margin-left: auto;
      }
    }
  }
}

.table-title {
  width: 10%;
  margin: 0 auto;
  font-weight: bold;
}
</style>